# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import csv
from collections import namedtuple
import datetime
import logging
import os
import sys

from absl import app, flags
import pathlib

from coin.support.proto_log.logic.util import run_from_og_log_archive
from coin.proto.coin_order_gateway_pb2 import OrderGatewayLog

FLAGS = flags.FLAGS

CsvWriter = namedtuple('CsvWriter', ['writer', 'file'])


class DumpStrategy(object):
  def __init__(self, exchange, csv_root):
    self._exchange = exchange
    self._csv_root = csv_root
    pathlib.Path(csv_root).mkdir(parents=True, exist_ok=True)
    self._prepare_csv_file()

  def _prepare_csv_file(self):
    csv_file = open('%s/%s.csv' % (self._csv_root, 'balance'), 'w', encoding='utf-8')
    writer = csv.writer(csv_file, dialect='excel')
    csv_header = ['timestamp', 'wallet_balance']
    writer.writerow(csv_header)
    csv_file.flush()
    self._balance_csv_writer = CsvWriter(writer, csv_file)

    csv_file = open('%s/%s.csv' % (self._csv_root, 'position'), 'w', encoding='utf-8')
    writer = csv.writer(csv_file, dialect='excel')
    csv_header = ['timestamp', 'net_position']
    writer.writerow(csv_header)
    csv_file.flush()
    self._position_csv_writer = CsvWriter(writer, csv_file)

    csv_file = open('%s/%s.csv' % (self._csv_root, 'order_event'), 'w', encoding='utf-8')
    writer = csv.writer(csv_file, dialect='excel')
    csv_header = [
        'timestamp',
        'price',
        'qty',
        'side',
        'internal_order_id',
        'fill_price',
        'fill_qty',
        'order_event_type'
    ]
    writer.writerow(csv_header)
    csv_file.flush()
    self._order_event_csv_writer = CsvWriter(writer, csv_file)

  def on_log(self, timestamp, log):
    pb = OrderGatewayLog()
    try:
      pb.ParseFromString(log)
    except Exception:
      print(log)
      return
    if pb.type == OrderGatewayLog.ORDER_EVENT:
      event = pb.event
      timestamp = event.event_time
      price = event.order_price
      qty = event.order_qty
      side = event.order_side
      internal_order_id = event.internal_order_id
      event_type = event.type
      fill_price = event.fill_price
      fill_qty = event.fill_qty
      if price == 0:
        price = float('nan')
      if fill_price == 0:
        fill_price = float('nan')
      csv_line = [timestamp, price, qty, side, internal_order_id, fill_price, fill_qty, event_type]
      writer = self._order_event_csv_writer
    elif pb.type == OrderGatewayLog.BALANCE:
      balances = pb.balance
      if self._exchange == 'Bitflyer':
        balances = [balance for balance in balances.each_balance if balance.currency == 'JPY']
      else:
        balances = balances.each_balance
      assert len(balances) == 1
      balance = balances[0]
      if self._exchange == 'Okex':
        wallet_balance = balance.available
      else:
        wallet_balance = balance.total
      csv_line = [timestamp, wallet_balance]
      writer = self._balance_csv_writer
    elif pb.type == OrderGatewayLog.POSITION:
      positions = pb.position
      if self._exchange == 'Okex':
        positions = [max(positions.each_position, key=lambda position: position.symbol)]
      else:
        positions = positions.each_position
      assert len(positions) == 1
      position = positions[0]
      net_position = position.net_position
      csv_line = [timestamp, net_position]
      writer = self._position_csv_writer
    else:
      # Don't care, skip.
      return

    writer.writer.writerow(csv_line)
    writer.file.flush()

  def close_csvs(self):
    self._balance_csv_writer.file.flush()
    self._position_csv_writer.file.close()
    self._order_event_csv_writer.file.close()


def launch(exchange, trading_date, root_dir, machine, csv_root):
  start_time = trading_date  # UTC 0 hour
  hours = 24

  strategy = DumpStrategy(exchange, csv_root)
  run_from_og_log_archive(strategy.on_log,
                          'Futures',
                          exchange,
                          start_time,
                          start_time + datetime.timedelta(hours=hours),
                          root_dir,
                          machine,
                          owner=None)
  strategy.close_csvs()
  return 0


def main(_):
  exchange = FLAGS.exchange
  trading_date = FLAGS.trading_date
  root_dir = os.path.expanduser(FLAGS.root_dir)
  machine = FLAGS.machine
  csv_root = FLAGS.csv_root
  assert trading_date, '--trading_date must be specified.'
  assert machine, '--machine must be specified.'
  assert exchange, '--exchange must be specified.'
  assert csv_root, '--csv_root must be specified.'
  print('Running for %s %s %s ...' % (trading_date, machine, exchange))
  sys.stdout.flush()
  trading_date = datetime.datetime.strptime(trading_date, '%Y%m%d')
  launch(exchange, trading_date, root_dir, machine, csv_root)


if __name__ == '__main__':
  flags.DEFINE_string('exchange', None, 'Exchange name.')

  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')

  flags.DEFINE_string('root_dir', '~/data/og_proto_log', 'root_dir.')

  flags.DEFINE_string('machine', None, 'Instance machine name.')

  flags.DEFINE_string('csv_root', None, 'Output csv files root directory.')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
